# Copyright (C) Kumo inc. and its affiliates.
# Author: Jeff.li lijippy@163.com
# All rights reserved.
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https:#www.gnu.org/licenses/>.
#

# Download Spark source code and copy the Spark connect proto files to Pollux
# directory.
set(SPARK_VERSION 3.5.1)
set(SPARK_SHA256_CHECKSUM
    SHA256=3ca674d1c7cc31e859ab45bf0e306e9b3329d78d142e75ff2ef162dd55b11383)
set(SPARK_SOURCE_URL
        "https://gitee.com/kumo-pub/spark/repository/archive/v${SPARK_VERSION}.tar.gz"
)

FetchContent_Declare(
  Spark
  URL ${SPARK_SOURCE_URL}
  URL_HASH ${SPARK_SHA256_CHECKSUM})
FetchContent_MakeAvailable(Spark)

set(SPARK_CONNECT_PROTO_BASE_DIR
    "${spark_SOURCE_DIR}/connector/connect/common/src/main/protobuf")
set(SPARK_CONNECT_PROTO_DIR "${SPARK_CONNECT_PROTO_BASE_DIR}/spark/connect")

# Set up Spark connect.
file(
  GLOB PROTO_FILES
  RELATIVE ${PROJECT_SOURCE_DIR}
  CONFIGURE_DEPENDS ${SPARK_CONNECT_PROTO_DIR}/*.proto)
foreach(PROTO ${PROTO_FILES})
  get_filename_component(PROTO_NAME ${PROTO} NAME_WE)
  list(APPEND PROTO_SRCS
       "${PROJECT_BINARY_DIR}/spark/connect/${PROTO_NAME}.pb.cc")
  list(APPEND PROTO_HDRS
       "${PROJECT_BINARY_DIR}/spark/connect/${PROTO_NAME}.pb.h")
  list(APPEND GRPC_SRCS
       "${PROJECT_BINARY_DIR}/spark/connect/${PROTO_NAME}.grpc.pb.cc")
  list(APPEND GRPC_HDRS
       "${PROJECT_BINARY_DIR}/spark/connect/${PROTO_NAME}.grpc.pb.h")
  list(APPEND PROTO_FILES_FULL "${SPARK_CONNECT_PROTO_DIR}/${PROTO_NAME}.proto")
endforeach()

set(PROTO_OUTPUT_FILES ${PROTO_HDRS} ${PROTO_SRCS})
set_source_files_properties(${PROTO_OUTPUT_FILES} PROPERTIES GENERATED TRUE)

set(GRPC_OUTPUT_FILES ${GRPC_HDRS} ${GRPC_SRCS})
set_source_files_properties(${GRPC_OUTPUT_FILES} PROPERTIES GENERATED TRUE)

# Ensure that the option --proto_path is not given an empty argument.
foreach(PROTO_PATH ${SPARK_CONNECT_PROTO_BASE_DIR} ${Protobuf_INCLUDE_DIRS})
  list(APPEND PROTO_PATH_ARGS --proto_path=${PROTO_PATH})
endforeach()

# Generate Spark connect hearders and sources.
add_custom_command(
  OUTPUT ${PROTO_OUTPUT_FILES}
  COMMAND protobuf::protoc ${PROTO_PATH_ARGS} --cpp_out ${PROJECT_BINARY_DIR}
          ${PROTO_FILES_FULL}
  DEPENDS protobuf::protoc
  COMMENT "Running PROTO compiler"
  VERBATIM)

# Generate grpc headers and sources.
add_custom_command(
  OUTPUT ${GRPC_OUTPUT_FILES}
  COMMAND
    protobuf::protoc ${PROTO_PATH_ARGS} --grpc_out=${PROJECT_BINARY_DIR}
    --plugin=protoc-gen-grpc=$<TARGET_FILE:gRPC::grpc_cpp_plugin>
    ${PROTO_FILES_FULL}
  DEPENDS protobuf::protoc
  COMMENT "Running gRPC C++ protocol buffer compiler"
  VERBATIM)

add_library(pollux_spark_query_runner ${PROTO_SRCS} ${GRPC_SRCS}
                                     SparkQueryRunner.cpp)

target_include_directories(pollux_spark_query_runner
                           PUBLIC ${PROJECT_BINARY_DIR}/spark/connect/)

target_link_libraries(
  pollux_spark_query_runner
  pollux_dwio_arrow_parquet_writer
  pollux_common_fuzzer_util
  pollux_fuzzer_util
  pollux_exec_test_lib
  pollux_arrow_bridge
  gRPC::grpc++
  gRPC::grpc
  protobuf::libprotobuf
        Arrow::arrow_static)

add_executable(spark_aggregation_fuzzer_test SparkAggregationFuzzerTest.cpp)

target_link_libraries(
  spark_aggregation_fuzzer_test
  pollux_aggregation_fuzzer
  pollux_aggregation_fuzzer_base
  pollux_functions_spark_aggregates
  pollux_spark_query_runner
  pollux_fuzzer_util
  pollux_window
  pollux_vector_test_lib
  GTest::gtest
  GTest::gtest_main)

if(POLLUX_BUILD_TESTING)
  add_subdirectory(tests)
endif()
